package com.senseauto.surfacerecord;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import android.util.Log;
import android.view.Surface;

import java.nio.ByteBuffer;

public class H264Encoder extends Thread{
    String TAG = "H264Encoder";
    private MediaCodec mediaCodec;
    private MediaMuxer mediaMuxer;
    private Surface recordSurface;
    private int trackIndex = -1;
    private boolean muxerStarted = false;

    private void setupMediaCodec(int width, int height, String path) {
        MediaFormat format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height);

        format.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);
        format.setInteger(MediaFormat.KEY_BIT_RATE, 12000000); // 比特率
        format.setInteger(MediaFormat.KEY_FRAME_RATE, 30);    // 帧率
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1); // 每秒一个关键帧

        try {
            mediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
            mediaCodec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);

            // 获取编码器的输入 Surface
            recordSurface = mediaCodec.createInputSurface();
            mediaCodec.start();

            // 初始化 MediaMuxer
            mediaMuxer = new MediaMuxer(path, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
        } catch (Exception e) {
            e.printStackTrace();
        }

    }

    private boolean isRecording = false;

    public Surface getRecordSurface(){
        return recordSurface;
    }

    public void stopRecord(){
        isRecording = false;
    }
    public void startRecording(int videoWidth, int videoHeight, String outputPath) {
        Log.e(TAG , "startRecording ->videoWidth:"+videoWidth+" videoHeight="+videoHeight+" outputPath="+outputPath);
        // 初始化 MediaCodec 和 MediaMuxer
        setupMediaCodec(videoWidth, videoHeight, outputPath);

        // 开启录制标志
        isRecording = true;

        // 启动一个线程处理编码器输出
        new Thread(() -> {
            while (isRecording) {
                drainEncoder();
            }

            // 在录制结束时确保所有缓冲区被处理
            drainEncoder();
            releaseResources();
        }).start();
    }


    private void drainEncoder() {
        MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
        int outputBufferIndex;

        while ((outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, 10000)) >= 0) {
            ByteBuffer outputBuffer = mediaCodec.getOutputBuffer(outputBufferIndex);

            if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0) {
                bufferInfo.size = 0;
            }


            if (bufferInfo.size != 0) {
                if (!muxerStarted) {
                    // 配置 MediaMuxer
                    MediaFormat format = mediaCodec.getOutputFormat();
                    trackIndex = mediaMuxer.addTrack(format);
                    mediaMuxer.start();
                    Log.e(TAG, "mediaMuxer.start();");
                    muxerStarted = true;
                }

                // 写入数据到 MediaMuxer
                outputBuffer.position(bufferInfo.offset);
                outputBuffer.limit(bufferInfo.offset + bufferInfo.size);

                mediaMuxer.writeSampleData(trackIndex, outputBuffer, bufferInfo);
            }

            mediaCodec.releaseOutputBuffer(outputBufferIndex, false);

            // 检查结束标志
            if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
                break;
            }
        }
    }

    private void releaseResources() {
        if (mediaCodec != null) {
            mediaCodec.stop();
            mediaCodec.release();
        }
        if (mediaMuxer != null) {
            if (muxerStarted) {
                mediaMuxer.stop();
            }
            mediaMuxer.release();
        }
    }
}